'''
This module estimates the RPKM from the bed(wig) files
by excluding the low coverage regions, which are likely to be 
exons. 
Created on Aug 22, 2009
@author: yanghoch
'''
import sys
import os
import getopt
from TophatOutFormat import WigLineInfo
from search_a_large_sorted_file import SortedFileSearcher

def getNextValidLine( filep, expColNo ):
    '''get the line with expected function'''
    delimiter = '\t'
    while True:
        line = ''
        line = filep.readline()
        if not line:
            return ""
        elif len( line.split( delimiter ) ) == expColNo:
            return line

def estimate_RPK_from_WIG( iFileN, range, coverageThreshold, readL ):
    ''' Given a wig file, the function estimates the number of reads
    mapped to the "exon" regions in the given range and normalized to
    1000 base pair. To get RPKM, the returned number should be 
    further normalized with the mapped reads number.     
    '''
    totalExonsLength = 0
    totalCoverage = 0

    posColId = 1 # wig format has its position in the 2nd column
    expColNo = 4
    s = SortedFileSearcher( iFileN );
    delimiter = '\t'
    getNextValidLine( s.fp, expColNo ) # skip the comment line
    s.move2the_begining_of_a_line( s.fp.tell() - 1 )
    s.locatedFp( min( range ), posColId, delimiter )
    while True:
        line = getNextValidLine( s.fp, expColNo )
        if not line: break
        wigline = WigLineInfo( line )
        if wigline.end > max( range ):
            wigline.end = max( range )
        if( wigline.coverage > coverageThreshold ):
            exonL = ( wigline.end - wigline.start )
            totalExonsLength += exonL
            totalCoverage += ( wigline.coverage * exonL )
        if wigline.end == max( range ): break
    if( totalExonsLength > 0 ):
        return( 1000 * totalCoverage / readL / totalExonsLength )
    else:
        return 0

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__
    print "Three arguments are required."
    print "The 1st argument is the wig file name."
    print "The 2nd and 3rd arguments are the range."

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();
    if ( len( args ) == 3 ):
        coverageThreshold = 0
        readL = 100
        iFileN = args[0]
        range = [int( args[1] ), int ( args[2] )]
        PRK = estimate_RPK_from_WIG( iFileN, range, coverageThreshold, readL )
        print "The estimated PRK is ", PRK
    else:
        usage()
if __name__ == '__main__':
    main()
